# #模型推理
# 首先，通过load_checkpoint与load_param_into_net将训练好的模型参数加载入新实例化的模型中。
import re

from text2text.model_training import src_vocab_size, d_model, n_heads, d_ff, n_layers, src_vocab, \
    trg_vocab, src_pad_idx, trg_pad_idx, trg_vocab_size
from transformers_network.decoder import Decoder
from transformers_network.encoder import Encoder
from transformers_network.transformer import Transformer

from mindspore import load_checkpoint, load_param_into_net, Tensor, ops
from mindspore import dtype as mstype

encoder = Encoder(src_vocab_size, d_model, n_heads, d_ff, n_layers, dropout_p=0.1)
decoder = Decoder(trg_vocab_size, d_model, n_heads, d_ff, n_layers, dropout_p=0.1)
new_model = Transformer(encoder, decoder)

ckpt_file_name = r"E:\transformer_for_text2text\text2text\transformer_en2cn.ckpt"
param_dict = load_checkpoint(ckpt_file_name)
load_param_into_net(new_model, param_dict)

"""
推理过程中无需对模型参数进行更新，所以这里model.set_train(False)。
我们输入一个英文语句，期望可以返回翻译好的中文语句。
首先通过Encoder提取英文序列中的特征信息，并将这些特征信息传输至Decoder。
Decoder最开始的输入为起始占位符 < bos >，每次会根据输入预测下一个出现的单词，并对输入进行更新，直到预测出终止占位符 < eos > 。
"""


def inference(sentence, max_len=41, target_len_max: int = None):
    """模型推理：输入一个英文句子，输出翻译后的中文句子
    enc_inputs: [batch_size(1), src_len]
    """
    new_model.set_train(False)

    # 对输入句子进行分词
    if isinstance(sentence, str):
        tokens = [tok.lower() for tok in re.findall(r'\w+|[^\w\s]', sentence.rstrip())]
    else:
        tokens = [token.lower() for token in sentence]

    n = len(tokens) # 输入句子的“单词”的个数

    # 补充起始、终止占位符，统一序列长度
    if len(tokens) > max_len - 2:
        src_len = max_len
        tokens = ['<bos>'] + tokens[:max_len - 2] + ['<eos>']
    else:
        src_len = len(tokens) + 2
        # tokens = ['<bos>'] + tokens + ['<eos>'] + ['<pad>'] * (max_len - src_len)
        tokens = ['<bos>'] + tokens + ['<pad>'] * (max_len - src_len) + ['<eos>']

    # 将源语言的单词转换为数字索引，并进一步转换为tensor
    # enc_inputs: [1, src_len]
    indexes = src_vocab.encode(tokens)
    enc_inputs = Tensor(indexes, mstype.float32).expand_dims(0)

    # 将输入送入encoder，获取信息
    enc_outputs, _ = new_model.encoder(enc_inputs, src_pad_idx)

    dec_inputs = Tensor([[src_vocab.bos_idx]], mstype.float32)

    # 初始化decoder输入，此时仅有句首占位符<pad>，即 dec_inputs: [1, 1]
    if target_len_max is None:
        target_len_max = n + 3
    for _ in range(target_len_max):
        dec_outputs, _, _ = new_model.decoder(dec_inputs, enc_inputs, enc_outputs, src_pad_idx, trg_pad_idx)
        dec_logits = dec_outputs.view((-1, dec_outputs.shape[-1]))

        # 找到下一个词的概率分布，并输出预测
        dec_logits = dec_logits[-1, :]
        pred = dec_logits.argmax(axis=0).expand_dims(0).expand_dims(0)
        pred = pred.astype(mstype.float32)
        # 更新dec_inputs
        dec_inputs = ops.concat((dec_inputs, pred), axis=1)
        # 如果出现<eos>，则终止循环
        if int(pred.asnumpy()[0, 0]) == trg_vocab.eos_idx:
            break
    # 最后需要将数字索引转换为英文单词
    trg_indexes = [int(i) for i in dec_inputs.view(-1).asnumpy()]
    # 如果有终止符号<eos>, 则去掉(如果预测序列存在<eos>则其一定是最后一个单词, 因为在预测下一个单词的过程中，只要预测出了<eos>，就退出循环了)
    if trg_indexes[-1] == trg_vocab.eos_idx:
        trg_indexes = trg_indexes[1:-1]   # 去掉起始、终止符号
    else:
        trg_indexes = trg_indexes[1:]  # 去掉起始符号, 因为没有终止符号, 所以不用去掉

    trg_tokens = trg_vocab.decode(trg_indexes)

    return trg_tokens


src_sentence = "Welcome to my home."
pred_trg = inference(src_sentence)

print(f'src = {src_sentence}')
print(f"predicted trg = {pred_trg}")

while True:
    src_sentence = input("input src sentence: ")
    if src_sentence == 'exit':
        print("Program exit.")
        break
    pred_trg = inference(src_sentence)
    print(f"predicted trg = {pred_trg}")
